# -*- coding: utf-8 -*-
# @Time      : 2025-08-01
# @Author    : 2023090007@ahdy.edu.cn
# @File      : loghandler
# @Desc  分布式锁
import redis
import uuid
import time
from contextlib import contextmanager

class RedisDistributedLock:
    def __init__(self, redis_client: redis.Redis, lock_key: str, 
                 expire_seconds: int = 30, retry_times: int = 3, 
                 retry_delay: float = 0.1):
        """
        初始化分布式锁
        
        :param redis_client: Redis客户端实例
        :param lock_key: 锁的键名
        :param expire_seconds: 锁的过期时间，防止死锁
        :param retry_times: 获取锁失败时的重试次数
        :param retry_delay: 重试间隔时间(秒)
        """
        self.redis = redis_client
        self.lock_key = lock_key
        self.expire_seconds = expire_seconds
        self.retry_times = retry_times
        self.retry_delay = retry_delay
        self.lock_value = None  # 存储当前锁的值，用于释放锁时验证
        self.is_acquired = False  # 锁是否已获取

    def acquire(self) -> bool:
        """获取分布式锁"""
        # 生成唯一的锁值，用于释放锁时验证身份
        self.lock_value = str(uuid.uuid4())
        
        for _ in range(self.retry_times):
            # 使用SET NX EX命令实现原子操作
            # NX: 只有键不存在时才设置成功
            # EX: 设置过期时间
            result = self.redis.set(
                self.lock_key, 
                self.lock_value, 
                nx=True, 
                ex=self.expire_seconds
            )
            
            if result:
                self.is_acquired = True
                return True
            
            # 重试前等待一段时间
            time.sleep(self.retry_delay)
        
        # 重试次数耗尽仍未获取到锁
        self.lock_value = None
        return False

    def release(self) -> bool:
        """释放分布式锁"""
        if not self.is_acquired or not self.lock_value:
            return False
        
        # 使用Lua脚本保证释放锁的原子性
        # 确保只有持有锁的客户端才能释放锁
        script = """
        if redis.call('get', KEYS[1]) == ARGV[1] then
            return redis.call('del', KEYS[1])
        else
            return 0
        end
        """
        
        # 替换脚本中的ARGV[1]为实际的锁值
        script = script.replace('ARGV[1]', f'"{self.lock_value}"')
        result = self.redis.eval(script, 1, self.lock_key)
        
        self.is_acquired = False
        self.lock_value = None
        return result == 1

    @contextmanager
    def lock(self):
        """上下文管理器，用于with语句"""
        try:
            acquired = self.acquire()
            yield acquired
        finally:
            if acquired:
                self.release()

# 使用示例
if __name__ == "__main__":
    # 初始化Redis客户端
    redis_client = redis.Redis(
        host='localhost',
        port=6379,
        db=0,
        decode_responses=True
    )
    
    # 创建分布式锁实例
    lock = RedisDistributedLock(
        redis_client=redis_client,
        lock_key="distributed_lock:example",
        expire_seconds=10,
        retry_times=5,
        retry_delay=0.2
    )
    
    # 方法1: 直接使用acquire和release
    if lock.acquire():
        try:
            print("成功获取锁，执行临界区操作...")
            # 模拟业务操作
            time.sleep(2)
        finally:
            lock.release()
            print("锁已释放")
    else:
        print("获取锁失败")
    
    # 方法2: 使用上下文管理器(with语句)
    with lock.lock() as acquired:
        if acquired:
            print("通过上下文管理器获取锁，执行临界区操作...")
            # 模拟业务操作
            time.sleep(2)
        else:
            print("通过上下文管理器获取锁失败")
